"""基础量化模型抽象类"""
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, Union
import pandas as pd
import numpy as np
from datetime import datetime
import os
import pickle


class BaseQuantModel(ABC):
    """量化模型基类，定义统一接口"""
    
    def __init__(self, model_name: str, task_type: str = "classification"):
        """
        初始化基础模型
        
        Args:
            model_name: 模型名称
            task_type: 任务类型，"classification"（分类）或 "regression"（回归）
        """
        self.model_name = model_name
        self.task_type = task_type
        self.model = None
        self.is_trained = False
        self.feature_columns = None
        self.model_params = {}
        
    @abstractmethod
    def _build_model(self, **kwargs):
        """构建模型（子类实现）"""
        pass
    
    @abstractmethod
    def _train_model(self, X_train: pd.DataFrame, y_train: pd.Series, **kwargs):
        """训练模型（子类实现）"""
        pass
    
    def train(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.Series] = None,
        **kwargs
    ) -> Dict:
        """
        训练模型
        
        Args:
            X_train: 训练特征数据
            y_train: 训练标签数据
            X_val: 验证特征数据（可选）
            y_val: 验证标签数据（可选）
            **kwargs: 其他训练参数
            
        Returns:
            训练结果字典
        """
        if self.model is None:
            self._build_model(**kwargs)
        
        # 保存特征列名
        self.feature_columns = list(X_train.columns)
        
        # 训练模型
        train_result = self._train_model(X_train, y_train, X_val, y_val, **kwargs)
        
        self.is_trained = True
        return train_result
    
    @abstractmethod
    def _predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测概率（分类任务）或直接预测值（回归任务）"""
        pass
    
    def predict(self, X: pd.DataFrame, return_proba: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        """
        预测
        
        Args:
            X: 特征数据
            return_proba: 是否返回概率（仅分类任务）
            
        Returns:
            预测结果，如果是分类任务且return_proba=True，返回(预测类别, 概率)
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，请先调用train()方法")
        
        if self.task_type == "classification":
            if return_proba:
                proba = self._predict_proba(X)
                predictions = np.argmax(proba, axis=1) if proba.ndim > 1 else (proba > 0.5).astype(int)
                return predictions, proba
            else:
                proba = self._predict_proba(X)
                predictions = np.argmax(proba, axis=1) if proba.ndim > 1 else (proba > 0.5).astype(int)
                return predictions
        else:
            return self._predict_proba(X)
    
    def backtest(
        self,
        X: pd.DataFrame,
        y: pd.Series,
        prices: pd.Series,
        initial_capital: float = 100000.0,
        commission_rate: float = 0.001,
        slippage_rate: float = 0.0005,
        position_limit: float = 1.0,
        stop_loss: Optional[float] = None,
        take_profit: Optional[float] = None,
        **kwargs
    ) -> Dict:
        """
        回测方法（walk-forward回测）
        
        Args:
            X: 特征数据
            y: 真实标签（用于计算胜率）
            prices: 价格序列（用于计算收益）
            initial_capital: 初始资金
            commission_rate: 交易手续费率
            slippage_rate: 滑点率
            position_limit: 最大仓位比例
            stop_loss: 止损比例（可选）
            take_profit: 止盈比例（可选）
            **kwargs: 其他回测参数
            
        Returns:
            回测结果字典，包含各项评估指标
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，请先调用train()方法")
        
        # 获取预测结果
        predictions = self.predict(X)
        
        # 如果是分类任务，将预测转换为信号（1买入，0卖出/持有，-1卖出）
        if self.task_type == "classification":
            signals = predictions.copy()
            signals[signals == 0] = -1  # 下跌预测转为卖出信号
        else:
            # 回归任务：正收益预测为买入，负收益预测为卖出
            signals = np.where(predictions > 0, 1, -1)
        
        # 初始化回测变量
        capital = initial_capital
        positions = 0  # 持仓数量
        entry_price = None
        equity_curve = [initial_capital]
        trades = []
        
        # 计算每日收益率
        returns = prices.pct_change().fillna(0).values
        
        for i in range(len(signals)):
            current_price = prices.iloc[i]
            signal = signals[i]
            
            # 计算当前权益
            current_equity = capital + positions * current_price
            
            # 处理止损/止盈
            if positions > 0 and entry_price is not None:
                pnl_rate = (current_price - entry_price) / entry_price
                if stop_loss is not None and pnl_rate <= -stop_loss:
                    signal = -1  # 触发止损
                elif take_profit is not None and pnl_rate >= take_profit:
                    signal = -1  # 触发止盈
            
            # 执行交易信号
            if signal == 1 and positions == 0:  # 买入
                max_position_value = current_equity * position_limit
                positions = max_position_value / current_price
                cost = positions * current_price * (1 + commission_rate + slippage_rate)
                capital -= cost
                entry_price = current_price
                trades.append({
                    'date': X.index[i] if hasattr(X.index, '__getitem__') else i,
                    'action': 'buy',
                    'price': current_price,
                    'shares': positions,
                    'cost': cost
                })
            elif signal == -1 and positions > 0:  # 卖出
                revenue = positions * current_price * (1 - commission_rate - slippage_rate)
                capital += revenue
                trades.append({
                    'date': X.index[i] if hasattr(X.index, '__getitem__') else i,
                    'action': 'sell',
                    'price': current_price,
                    'shares': positions,
                    'revenue': revenue
                })
                positions = 0
                entry_price = None
            
            # 更新权益曲线
            current_equity = capital + positions * current_price
            equity_curve.append(current_equity)
        
        # 计算最后持仓价值
        if positions > 0:
            final_price = prices.iloc[-1]
            capital += positions * final_price * (1 - commission_rate - slippage_rate)
        
        equity_curve = np.array(equity_curve)
        
        # 计算评估指标
        total_return = (capital - initial_capital) / initial_capital
        
        # 计算每日收益率
        equity_returns = np.diff(equity_curve) / equity_curve[:-1]
        equity_returns = equity_returns[~np.isnan(equity_returns)]
        
        if len(equity_returns) == 0:
            annualized_return = 0
            volatility = 0
            sharpe_ratio = 0
            max_drawdown = 0
            information_ratio = 0
        else:
            # 年化收益率
            trading_days = len(equity_returns)
            if trading_days > 0:
                annualized_return = (1 + total_return) ** (252 / trading_days) - 1
            else:
                annualized_return = 0
            
            # 年化波动率
            volatility = np.std(equity_returns) * np.sqrt(252)
            
            # 夏普比率（假设无风险利率为0）
            sharpe_ratio = annualized_return / volatility if volatility > 0 else 0
            
            # 最大回撤
            cummax = np.maximum.accumulate(equity_curve)
            drawdown = (equity_curve - cummax) / cummax
            max_drawdown = np.min(drawdown)
            
            # 信息比率（相对于基准）
            if len(returns) == len(equity_returns):
                excess_returns = equity_returns - returns
                information_ratio = np.mean(excess_returns) / np.std(excess_returns) * np.sqrt(252) if np.std(excess_returns) > 0 else 0
            else:
                information_ratio = 0
        
        # 胜率（仅分类任务）
        if self.task_type == "classification" and len(predictions) == len(y):
            correct_predictions = (predictions == y.values).sum()
            win_rate = correct_predictions / len(predictions) if len(predictions) > 0 else 0
        else:
            win_rate = None
        
        return {
            'model_name': self.model_name,
            'total_return': total_return,
            'annualized_return': annualized_return,
            'volatility': volatility,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'win_rate': win_rate,
            'information_ratio': information_ratio,
            'final_capital': capital,
            'total_trades': len(trades),
            'equity_curve': equity_curve,
            'trades': trades
        }
    
    def save(self, filepath: str):
        """保存模型到文件"""
        if not self.is_trained:
            raise ValueError("模型尚未训练，无法保存")
        
        model_data = {
            'model': self.model,
            'model_name': self.model_name,
            'task_type': self.task_type,
            'feature_columns': self.feature_columns,
            'model_params': self.model_params,
            'is_trained': self.is_trained
        }
        
        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else '.', exist_ok=True)
        
        with open(filepath, 'wb') as f:
            pickle.dump(model_data, f)
    
    def load(self, filepath: str):
        """从文件加载模型"""
        with open(filepath, 'rb') as f:
            model_data = pickle.load(f)
        
        self.model = model_data['model']
        self.model_name = model_data['model_name']
        self.task_type = model_data['task_type']
        self.feature_columns = model_data['feature_columns']
        self.model_params = model_data.get('model_params', {})
        self.is_trained = model_data['is_trained']

